import os
import argparse

"""
extract images from yuv videos
"""

parser = argparse.ArgumentParser()
parser.add_argument("-f", "--seq_folder", type = str)

parser.add_argument("-o", "--save_folder", type = str, default = "./compressed")
parser.add_argument("--format", type = str, default = "png")

args = parser.parse_args()

seq_folder = args.seq_folder

save_folder = args.save_folder
format = args.format


if not os.path.exists(save_folder):
    os.makedirs(save_folder)


for seq in os.listdir(seq_folder):
    if seq.endswith(".yuv"):
        input = os.path.join(seq_folder, seq)
        output_folder = os.path.join(save_folder, seq[:-4])

        if not os.path.exists(output_folder):
            os.makedirs(output_folder)

        options = [
            "ffmpeg",
            "-i", input,
            "-f", "image2",
            f"{output_folder}/im%04d.{format}"
        ]

        cmd = ""
        for opt in options:
            cmd += (opt + " ").replace("\\", "/")

        print(cmd)
        os.system(cmd)

